import torch
import torchvision
from torch import nn
from torch.nn import ReLU
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision.utils import make_grid

# input = torch.tensor([
#     [1, -0.5],
#     [-1, 3]
# ])
# input = torch.reshape(input, (-1, 1, 2, 2))
'''
1.加载CIFAR10数据集，并通过DataLoader批量加载图像数据。
2.定义一个包含ReLU层的神经网络模块，用于对输入图像进行最大池化操作。
3.使用TensorBoard记录输入图像和经过非线性激活后的输出图像。
4.遍历数据加载器，对每批数据应用最大池化操作，并将结果写入日志文件。
'''
dataset = torchvision.datasets.CIFAR10("../../data", train=False, transform=torchvision.transforms.ToTensor(),
                                       download=True)
dataLoader = DataLoader(dataset, batch_size=64)

class Module(nn.Module):
    def __init__(self):
        super(Module, self).__init__()
        self.relu1 = ReLU()
    def forward(self, x):
        x = self.relu1(x)
        return x

module = Module()
# input = module(input)
# print(input)
step = 0
writer = SummaryWriter("../../logs_relu")
for data in dataLoader:
    imgs, targets = data
    grid_imgs = make_grid(imgs)
    writer.add_image("input", grid_imgs, global_step=step)
    output = module(imgs)
    grid_output = make_grid(output)
    writer.add_image("output", grid_output, global_step=step)
    step += 1
writer.close()